﻿// my_reduce.cpp
#include <torch/extension.h>  // 包含了必要的头文件，包括ATen、pybind11等

// 前向传播函数
torch::Tensor my_reduce_fwd(const torch::Tensor& y)
{
    // y 可以是任意 shape，这里举个最简单的：sum 规约
    return y.sum();  // 你可以替换成自己手写的 SIMD / OpenMP / CUDA 实现
}

// 反向传播函数：把上游梯度广播成 y 的形状
torch::Tensor my_reduce_bwd(const torch::Tensor& grad_out, const torch::Tensor& y) { return grad_out.expand_as(y); }

// 使用 TorchScript 自定义操作的方式注册正反两个函数，并关联自动微分
// 这是新版本推荐的方式，替代旧的 PYBIND11_MODULE 方式:cite[6]:cite[7]
TORCH_LIBRARY(my_ops, m)
{
    // 注册前向函数 "my_reduce"
    m.def("my_reduce", my_reduce_fwd);
    // 注册反向函数 "my_reduce_backward"
    m.def("my_reduce_backward", my_reduce_bwd);

    // 如果你的操作需要支持自动微分，通常不需要像上面那样手动分开注册正反向
    // 更常见的做法是定义一个继承自 torch::autograd::Function 的类
    // 但 TORCH_LIBRARY 方式更底层，适用于与 PyTorch 核心深度集成
}

// 注意：在新版本的 PyTorch C++ 扩展实践中，如果使用了 TORCH_LIBRARY，
// 通常就不再需要（或者说不应该同时使用）PYBIND11_MODULE 来暴露同一个函数给 Python 了。
// 你应该选择一种方式。对于希望深度集成、高效且能被 TorchScript 识别的操作，
// 推荐使用 TORCH_LIBRARY。

// 如果你仍然需要传统的 Python 绑定（例如为了兼容性，或者需要被直接调用），
// 可以保留 PYBIND11_MODULE，但要注意避免重复定义和链接冲突。
// 但更现代的做法是让 Python 通过 torch.ops.your_namespace.your_operator 来调用
// TORCH_LIBRARY 注册的操作:cite[6]。
#if 0
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    // 如果你的构建系统（如 setup.py 或 CMakeLists.txt）定义了 TORCH_EXTENSION_NAME，
    // 例如为 'my_reduce'，那么这个模块名就是 'my_reduce'
    m.def("forward", &my_reduce_fwd, "custom reduce forward");
    m.def("backward", &my_reduce_bwd, "custom reduce backward");
    // 注意：通过 pybind11 这样直接暴露的 'forward' 和 'backward' 函数，
    // 不会自动与 PyTorch 的 autograd 系统关联。它们只是普通的函数。
    // 通常，更常见的做法是定义一个继承 torch.autograd.Function 的 Python 类，
    // 在其 forward 和 backward 方法中分别调用 C++ 的 forward 和 backward 函数:cite[1]:cite[7]。
}
#endif